简单好用！北大、普林斯顿联合提出即插即用的大语言模型加速方法
最近，大语言模型（LLM）生成过程的加速技术，例如投机解码、Medusa（美杜莎）等，都带来了令人印象深刻的速度提升。这些方法通常依赖于将 LLM 与一个小型的草稿模型配对。小型的草稿模型试图在每个解码步骤中以更低的延迟预测多个草稿 token，并让 LLM 并行验证它们，从而减少了 LLM 的解码步数。
然而，获得高质量的草稿 token 仍然是一门艺术：它必须在 LM 更少量的参数和更强大的预测能力之间找到平衡，同时草稿模型必须匹配 LLM 的词汇表；此外，它还应该便于集成到用于部署 LLM 的分布式系统中。
为了应对这些挑战，Medusa 引入了一个有效的微调方法来获得草稿模型，它基于 LLM 本身微调了多个头（Medusa head）来预测草稿 token。然而，对额外微调的要求仍然有诸多不便。
这就引出了一个问题——能否设计一个即插即用的加速方法？一个无需训练或微调新模型即可提供快速生成草稿 token 的方法？
为了实现这个目标，作者引入了 REST: Retrieval-Based Speculative Decoding，一种基于检索的投机解码方法。REST 不使用额外的草稿模型，而是利用数据存储来根据部分输入来检索草稿 token。通过在 CPU 上进行检索，REST 避免了额外的 GPU 负载。由于不需要联合训练或微调，REST 可以立即加速任何预训练语言模型。
在这篇文章中，研究人员将探讨 REST 的灵感和内部工作机制，揭示它如何在没有额外训练模型的情况下实现令人印象深刻的速度提升，通过清晰的解释和生动的例子来展示 REST 的能力。
收录会议：
博客地址：
代码地址：
首先，来观察一下 Python 代码的生成过程。例如，当模型生成 “for i” 时，熟悉 Python 的程序员可以高度预测接下来的几个令牌很可能是 “in range(”。
然而，当使用像 CodeLlama 这样的 LLM 时，它必须依次生成 "in"，"range"，和 "(" 这三个 token。这极大地降低代码生成的速度。那么有没有更高效的方法来一次性生成这些常见的短语呢？
答案是有的！这就是 REST 进入画面的地方。
如图所示，在推理过程中，语境被当做 query 来从数据存储中检索相关的语境，所有的后面紧跟着的序列，可以当做候选序列。
候选序列中每个的所有前缀都有可能充当语境后面的 token。然而，候选序列的数量可能是很庞大的，由于 GPU 显存大小和运算能力的限制无法一次性喂给 LLM 来验证。
因此，作者在候选序列的基础上建立了一个字典树（前缀树）Trie，Trie 中的每个结点到根节点就形成了候选序列中的一个前缀，而结点的权重即是该前缀在候选序列中出现的次数。
基于 Trie，可以轻易地选出出现次数最多的多个前缀，来当做草稿序列。然而，作者观察到这些草稿序列中，会经常出现重复的前缀，比如图中 in range 和 in range( 都有前缀 in range, in range( 和 in sorted 都有前缀 in，如果直接喂给大模型验证的话，会造成对相同前缀的重复计算，浪费了计算资源。
因此，研究人员观察到出现次数最多的前缀实际上是从权重最大的多个结点构成的子树（subtree）中获得的，通过对该子树进行广度优先遍历（BFS），可以得到一个伪句子（pseudo sequence），由字典树的特性可得草稿序列一定都是这个伪句子的子序列。通过专门的设计自注意力掩码，可以在 LLM 验证的时候还原伪句子中 token 原本的依赖关系（也称作树注意力）。
最后，如果草稿序列中有多个草稿 token 被 LLM 接受，那么 LLM 的推理速度就得到了提高。注意，这里作者采用的接受策略是当且仅当草稿 token 和 LLM 本应该生成的 token 相同时，才接受。因此通过 REST 来加速生成的内容和用原本 LLM 自回归生成的内容是完全一样的。
在单张 A6000 GPU 上，研究人员对 CodeLlama 7B 和 13B 在 HumanEval 上进行测试，所用的 datastore大小为 27GB，取得了 2.12~2.36 倍的加速效果；对 Vicuna 7B 和 13B 在 MT-Bench 上进行了测试，所用的 datastore 大小为 12GB，取得了 1.62~1.77 倍的加速效果。
研究人员在不同 Datastore 大小上进行了实验，可以观察到当 Datastore 增大时，加速效果也更强。此外，可以观察到用于检索的时间（包括建立字典树）是基本可以忽略的，这也说明了 REST 的高效率。
更多阅读
#投 稿 通 道#
让你的文字被更多人看到
如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。
总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。
📝 稿件基本要求：
• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注
• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题
• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道：
• 投稿邮箱：hr@paperweekly.site
• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿
△长按添加PaperWeekly小编
🔍
现在，在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧